from torch_geometric.data import Data
from torch_geometric.nn import GCNConv
import torch.nn.functional as F
import torch.optim as optim
import torch

bom_structure = {
    "Part1": {
        "SBB1": [],
        "SBB2": ["SBB3"],
        "SBB3": ["SBB4"]
    },
    "Part2": {
        "SBB3": ["SBB4"],
        "SBB2": ["SBB1"],
        "SBB4": []
    },
    "Part3": {
        "SBB1": [],
        "SBB3": ["SBB5"]
    }
}

# 从bom_structure中提取所有组件到components列表
components = [part for product, contents in bom_structure.items() for part, _ in contents.items()]
components.extend(
    [subpart for product, contents in bom_structure.items() for part, subparts in contents.items() for subpart in
     subparts])

# 确保components是唯一的并排序
components = sorted(set(components))

# 创建从组件名称到索引的映射
component_to_index = {comp: index for index, comp in enumerate(components)}

# 构建边列表
edge_list = []
for product, contents in bom_structure.items():
    for part, subparts in contents.items():
        for subpart in subparts:
            edge_list.append((part, subpart))

# 将边列表转换为PyTorch Geometric所需的格式
edge_index = torch.tensor([[component_to_index[source], component_to_index[target]] for source, target in edge_list],
                          dtype=torch.long).t().contiguous()

# 接下来，您可以继续构建特征矩阵、Data对象、定义GNN模型等步骤

# 步骤3: 构建特征矩阵
n_components = len(components)
node_features = torch.eye(n_components)  # 使用单位矩阵作为节点特征

# 构建Data对象
bom_data = Data(x=node_features, edge_index=edge_index)


# 定义GNN模型
class GNNRecommender(torch.nn.Module):
    def __init__(self, num_node_features, hidden_size, out_size):
        super(GNNRecommender, self).__init__()
        self.conv1 = GCNConv(num_node_features, hidden_size)
        self.conv2 = GCNConv(hidden_size, out_size)

    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        x = self.conv2(x, edge_index)
        return x


# 实例化模型
hidden_size = 8  # 隐藏层大小
out_size = n_components  # 输出层大小与组件数量相同
model = GNNRecommender(node_features.size(1), hidden_size, out_size)

# 步骤6: 训练GNN模型
# 随机初始化一个标签矩阵，仅作为示例
bom_data.y = torch.rand((n_components, out_size))
# 实例化模型和优化器
model = GNNRecommender(node_features.size(1), hidden_size, out_size)
optimizer = optim.Adam(model.parameters(), lr=0.01)
# 训练模型
model.train()
for epoch in range(500):
    print("Start Iteration:" + str(epoch))
    optimizer.zero_grad()
    out = model(bom_data.x, bom_data.edge_index)
    loss = F.mse_loss(out, bom_data.y)
    loss.backward()
    optimizer.step()


def get_subcomponent_recommendations(model, data, parent_component):
    model.eval()
    component_to_recommendations = {}
    with torch.no_grad():
        # 获取所有组件的嵌入
        component_embeddings = model(data.x, data.edge_index)

        # 找到父部件的索引
        parent_index = components.index(parent_component)

        # 计算父部件与所有组件的相似度
        similarities = F.cosine_similarity(component_embeddings, component_embeddings[parent_index].unsqueeze(0))

        # 获取与父部件直接相连的子部件的索引
        child_indices = data.edge_index[1][data.edge_index[0] == parent_index]

        # 获取子部件的相似度得分
        child_similarities = similarities[child_indices]

        # 对子部件的相似度进行排序
        sorted_indices = child_similarities.sort(descending=True).indices
        sorted_scores = child_similarities.sort(descending=True).values

        # 使用列表推导式获取推荐子部件及其相似度得分
        recommended_subcomponents = [(components[child_indices[i]], sorted_scores[j].item()) for j, i in
                                     enumerate(sorted_indices.tolist())]
        component_to_recommendations[parent_component] = recommended_subcomponents

    return component_to_recommendations

parent_component = "SBB3"  # 假设您想为“组件A”找到推荐的子部件
recommendations = get_subcomponent_recommendations(model, bom_data, parent_component)
for component, recs in recommendations.items():
    print(f"{component}的推荐子部件及相似度得分: {recs}")